#!/usr/bin/env python3
"""
Pydantic v2 兼容性修复脚本
自动修复项目中的 Pydantic v1 到 v2 的兼容性问题
"""

import os
import re
import sys
from pathlib import Path

def fix_pydantic_imports(file_path):
    """修复 Pydantic 导入语句"""
    with open(file_path, 'r', encoding='utf-8') as f:
        content = f.read()
    
    original_content = content
    
    # 替换导入语句
    content = re.sub(
        r'from pydantic import (.*)validator(.*)root_validator(.*)',
        r'from pydantic import \1field_validator\2model_validator\3',
        content
    )
    
    content = re.sub(
        r'from pydantic import (.*)root_validator(.*)validator(.*)',
        r'from pydantic import \1model_validator\2field_validator\3',
        content
    )
    
    # 单独的导入
    content = re.sub(r'validator,', 'field_validator,', content)
    content = re.sub(r'root_validator,', 'model_validator,', content)
    content = re.sub(r', validator', ', field_validator', content)
    content = re.sub(r', root_validator', ', model_validator', content)
    
    return content, content != original_content

def fix_validator_decorators(file_path):
    """修复验证器装饰器"""
    with open(file_path, 'r', encoding='utf-8') as f:
        content = f.read()
    
    original_content = content
    
    # 修复 @validator
    content = re.sub(
        r'@validator\(([^)]+)\)\s*\n\s*def (\w+)\(cls, v\):',
        r'@field_validator(\1)\n    @classmethod\n    def \2(cls, v):',
        content,
        flags=re.MULTILINE
    )
    
    # 修复 @root_validator
    content = re.sub(
        r'@root_validator\s*\n\s*def (\w+)\(cls, values\):',
        r'@model_validator(mode=\'after\')\n    def \1(self):',
        content,
        flags=re.MULTILINE
    )
    
    # 修复 @root_validator(pre=True)
    content = re.sub(
        r'@root_validator\(pre=True\)\s*\n\s*def (\w+)\(cls, values\):',
        r'@model_validator(mode=\'before\')\n    @classmethod\n    def \1(cls, values):',
        content,
        flags=re.MULTILINE
    )
    
    return content, content != original_content

def fix_validator_logic(file_path):
    """修复验证器内部逻辑"""
    with open(file_path, 'r', encoding='utf-8') as f:
        content = f.read()
    
    original_content = content
    
    # 在 @model_validator(mode='after') 中，将 values.get() 替换为 self.field
    # 这需要更复杂的逻辑，暂时跳过自动修复
    
    return content, content != original_content

def process_file(file_path):
    """处理单个文件"""
    print(f"处理文件: {file_path}")
    
    try:
        # 修复导入
        content, changed1 = fix_pydantic_imports(file_path)
        
        # 修复装饰器
        content, changed2 = fix_validator_decorators(file_path)
        
        # 修复逻辑
        content, changed3 = fix_validator_logic(file_path)
        
        if changed1 or changed2 or changed3:
            with open(file_path, 'w', encoding='utf-8') as f:
                f.write(content)
            print(f"  ✅ 已修复")
            return True
        else:
            print(f"  ⏭️ 无需修复")
            return False
            
    except Exception as e:
        print(f"  ❌ 修复失败: {e}")
        return False

def main():
    """主函数"""
    print("🔧 开始修复 Pydantic v2 兼容性问题...")
    
    # 查找所有 Python 文件
    backend_dir = Path(__file__).parent.parent / 'backend'
    python_files = list(backend_dir.rglob('*.py'))
    
    # 过滤出包含 Pydantic 的文件
    pydantic_files = []
    for file_path in python_files:
        try:
            with open(file_path, 'r', encoding='utf-8') as f:
                content = f.read()
                if 'pydantic' in content and ('validator' in content or 'root_validator' in content):
                    pydantic_files.append(file_path)
        except Exception:
            continue
    
    print(f"找到 {len(pydantic_files)} 个需要修复的文件")
    
    fixed_count = 0
    for file_path in pydantic_files:
        if process_file(file_path):
            fixed_count += 1
    
    print(f"\n🎉 修复完成！共修复 {fixed_count} 个文件")
    
    if fixed_count > 0:
        print("\n📝 修复说明:")
        print("1. @validator -> @field_validator")
        print("2. @root_validator -> @model_validator")
        print("3. 添加了 @classmethod 装饰器")
        print("4. 部分验证器逻辑可能需要手动调整")
        
        print("\n⚠️ 注意事项:")
        print("1. 请检查修复后的代码是否正确")
        print("2. @model_validator(mode='after') 中需要使用 self.field 而不是 values.get('field')")
        print("3. 返回值需要是 self 而不是 values")

if __name__ == "__main__":
    main()
